﻿/*
 * Description: application window helper.
 */
#include <windows.h>
#include <WtsApi32.h>
#include <iostream>
#include <unordered_set>
#include "Psapi.h"
#include "common.h"

namespace {

enum WindowHelperErrorCode : uint32_t {
    SUCCESS = 0,
    MALLOC_FAIL,
    GET_USERNAME_FAIL,
    COMBINE_PATH_FAIL,
    CREATE_FILE_FAIL,
    WRITE_FILE_FAIL
};

std::unordered_set<DWORD> taskBarApplications;
std::unordered_set<HWND> openedWindowApplications;
std::unordered_set<DWORD> filterWindowApplications;

std::unordered_set<std::string> windowWhiteList = {
    "WindowsTerminal.exe",
    "powershell.exe"
};

BOOL IsWindowVisibleOnDesktop(HWND hwnd)
{
    if (hwnd == nullptr || hwnd == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    RECT windowRect = { 0 };
    if (!GetWindowRect(hwnd, &windowRect)) {
        return FALSE;
    }

    if (!IsWindowVisible(hwnd)) {
        return FALSE;
    }

    DWORD dwStyle = GetWindowLong(hwnd, GWL_STYLE);
    if ((dwStyle & WS_DISABLED) || (dwStyle & WS_MINIMIZE) || (dwStyle & WS_CHILD)) {
        return FALSE;
    }

    HWND hwndParent = GetAncestor(hwnd, GA_ROOTOWNER);
    if (hwndParent != NULL) {
        DWORD exStyle = GetWindowLong(hwndParent, GWL_EXSTYLE);
        if ((exStyle & WS_EX_TOOLWINDOW) || (exStyle & WS_EX_NOACTIVATE) ||
            (exStyle & WS_EX_TOPMOST) || (exStyle == WS_EX_LAYERED)) {
            return FALSE;
        }
    }

    if (((windowRect.right - windowRect.left) <= 0) || ((windowRect.bottom - windowRect.top) <= 0)) {
        return FALSE;
    }
    return TRUE;
}

BOOL IsWindowInTaskBar(HWND hwnd)
{
    if (hwnd == nullptr || hwnd == INVALID_HANDLE_VALUE) {
        return FALSE;
    }

    if (!IsWindowVisible(hwnd)) {
        return FALSE;
    }
    WINDOWPLACEMENT wp;
    SecureZeroMemory(&wp, sizeof(WINDOWPLACEMENT));
    wp.length = sizeof(WINDOWPLACEMENT);
    if (!GetWindowPlacement(hwnd, &wp)) {
        return FALSE;
    }

    if (wp.showCmd != SW_SHOWMINIMIZED) {
        return FALSE;
    }
    return TRUE;
}

std::string ConvertWStringToString(const std::wstring& wString)
{
    std::string str = "";
    int len = WideCharToMultiByte(CP_ACP, 0, wString.c_str(), static_cast<int>(wString.size()),
        NULL, 0, NULL, NULL);
    if (len <= 0) {
        return str;
    }

    char* buffer = new char[len + 1];
    if (buffer == nullptr) {
        return str;
    }

    SecureZeroMemory(buffer, (len + 1) * sizeof(char));
    WideCharToMultiByte(CP_ACP, 0, wString.c_str(), static_cast<int>(wString.size()), buffer, len, NULL, NULL);
    buffer[len] = '\0';
    str.assign(buffer);
    delete[] buffer;

    return str;
}

std::string GetFullNameByProcessId(DWORD pid)
{
    std::string fullName = "";
    WCHAR tempProcName[MAX_PATH] = { 0 };
    HANDLE processHandle = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, pid);
    if (processHandle != NULL) {
        GetModuleFileNameExW(processHandle, NULL, tempProcName, MAX_PATH);
        CloseHandle(processHandle);
        processHandle = NULL;
        std::wstring procName = tempProcName;
        fullName = ConvertWStringToString(procName);
    }
    return fullName;
}

std::string GetProcessNameFromFullPath(const std::string& path)
{
    if (path.empty()) {
        return path;
    }
    std::string processName = path;
    std::string::size_type idx = path.rfind("\\");
    if (idx != path.npos) {
        processName = processName.substr(idx + 1);
    }
    return processName;
}

BOOL CALLBACK EnumWindowsProc(HWND hwnd, LPARAM lParam)
{
    DWORD pid;
    GetWindowThreadProcessId(hwnd, &pid);
    if (IsWindowVisibleOnDesktop(hwnd)) {
        openedWindowApplications.insert(hwnd);
    }
    if (IsWindowInTaskBar(hwnd)) {
        taskBarApplications.insert(pid);
    }
    return TRUE;
}

void FilterWindowsApplication(HWND hwnd, std::unordered_set<DWORD>& collector,
    std::unordered_set<std::string>& whiteList)
{
    DWORD pid;
    GetWindowThreadProcessId(hwnd, &pid);
    DWORD exStyle = GetWindowLong(hwnd, GWL_EXSTYLE);
    if (exStyle & MB_SERVICE_NOTIFICATION) {
        std::string inputFullName = GetFullNameByProcessId(pid);
        if (inputFullName.size() != 0) {
            std::string procName = GetProcessNameFromFullPath(inputFullName);
            if (whiteList.find(procName) != whiteList.end()) {
                collector.insert(pid);
            }
        }
    } else {
        collector.insert(pid);
    }
}
}  // namespace

int main()
{
    LPSTR username = (LPSTR) malloc(sizeof(TCHAR) * MAX_PATH);
    if (username == nullptr) {
        return static_cast<uint32_t>(WindowHelperErrorCode::MALLOC_FAIL);
    }

    SecureZeroMemory(username, sizeof(TCHAR) * MAX_PATH);
    DWORD size = sizeof(TCHAR) * MAX_PATH;
    if (!GetUserName(username, &size)) {
        std::cout << "Failed to get username." << std::endl;
        free(username);
        username = nullptr;
        return static_cast<uint32_t>(WindowHelperErrorCode::GET_USERNAME_FAIL);
    }

    char combinePath[MAX_PATH] = { 0 };
    if (sprintf_s(combinePath, MAX_PATH, APPLICATION_WINDOW_HELPER_FILE_PATH, username) < 0) {
        std::cout << "Error occurred, combine path failed." << std::endl;
        free(username);
        username = nullptr;
        return static_cast<uint32_t>(WindowHelperErrorCode::COMBINE_PATH_FAIL);
    }

    EnumWindows(EnumWindowsProc, NULL);
    if (openedWindowApplications.size() != 0) {
        for (auto it = openedWindowApplications.begin(); it != openedWindowApplications.end(); ++it) {
            FilterWindowsApplication(*it, filterWindowApplications, windowWhiteList);
        }
    }

    bool isApplicationInTaskBar = filterWindowApplications.size() > 0 || taskBarApplications.size() > 0;
    HANDLE hFile = CreateFile(combinePath, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS,
        FILE_ATTRIBUTE_NORMAL, NULL);
    if (hFile == nullptr || hFile == INVALID_HANDLE_VALUE) {
        free(username);
        username = nullptr;
        return static_cast<uint32_t>(WindowHelperErrorCode::CREATE_FILE_FAIL);
    }

    free(username);
    username = nullptr;

    std::string result = isApplicationInTaskBar ? "1" : "0";
    DWORD bytesWritten;
    if (!WriteFile(hFile, result.c_str(), result.size() + sizeof TCHAR, &bytesWritten, NULL)) {
        std::cout << "Failed to write to file." << std::endl;
        CloseHandle(hFile);
        return static_cast<uint32_t>(WindowHelperErrorCode::WRITE_FILE_FAIL);
    }
    CloseHandle(hFile);
    return static_cast<uint32_t>(WindowHelperErrorCode::SUCCESS);
}